package com.lmk.dsa.struct.tree;

import com.lmk.dsa.struct.commons.BstNode;

/**
 * 平衡二叉树
 * @param <K>
 * @param <V>
 * @author LaoMake
 * @since 1.0
 */
public class AVLTree<K extends Comparable<? super K>, V> extends BinarySearchTree<K, V>{

    // 对节点y进行向右旋转操作，返回旋转后新的根节点x
    //        y                              x
    //       / \                           /   \
    //      x   T4     向右旋转 (y)        z     y
    //     / \       - - - - - - - ->    / \   / \
    //    z   T3                       T1  T2 T3 T4
    //   / \
    // T1   T2
    // 即：x.right = y; y.left = T3
    private BstNode<K, V> rightRotate(BstNode<K, V> y){
        BstNode x = y.left;
        BstNode T3 = x.right;
        x.right = y;
        y.left = T3;

        // 更新y、x的高度
        y.height = Math.max(nodeHeight(y.left), nodeHeight(y.right)) + 1;
        x.height = Math.max(nodeHeight(x.left), nodeHeight(x.right)) + 1;
        return x;
    }

    // 对节点y进行向左旋转操作，返回旋转后新的根节点x
    //    y                             x
    //  /  \                          /   \
    // T1   x      向左旋转 (y)       y     z
    //     / \   - - - - - - - ->   / \   / \
    //   T2  z                     T1 T2 T3 T4
    //      / \
    //     T3 T4
    // 即：x.left = y; y.right = x.left
    private BstNode<K, V> leftRotate(BstNode<K, V> y){
        BstNode x = y.right;
        BstNode T2 = x.left;
        x.left = y;
        y.right = T2;

        // 更新y、x的高度
        y.height = Math.max(nodeHeight(y.left), nodeHeight(y.right)) + 1;
        x.height = Math.max(nodeHeight(x.left), nodeHeight(x.right)) + 1;
        return x;
    }

    /**
     * 维护节点平衡
     * @param node
     */
    private BstNode<K, V> balanceNode(BstNode<K, V> node){
        if(node == null)
            return null;

        BstNode<K, V> retNode = node;

        // 更新当前父节点的高度
        node.height = 1 + Math.max(nodeHeight(node.left), nodeHeight(node.right));

        // 获得当前父节点的平衡因子
        int balanceFactor = getBalanceFactor(node);

        // 在当前节点的左侧添加时产生不平衡
        if(balanceFactor > 1){
            if(getBalanceFactor(node.left) >= 0){
                // LL 左侧的左侧
                retNode = rightRotate(node);
            }else{
                // LR 左侧的右侧
                // 先对x左旋
                node.left = leftRotate(node.left);
                // 当前节点右旋
                retNode = rightRotate(node);
            }
        }

        // 在当前节点的右侧添加时产生不平衡
        if(balanceFactor < -1){
            if(getBalanceFactor(node.right) <= 0){
                // RR 左侧的左侧
                retNode = leftRotate(node);
            }else{
                // RL 左侧的右侧
                // 先对x右旋
                node.right = rightRotate(node.right);
                // 当前节点左旋
                retNode = leftRotate(node);
            }
        }
        return retNode;
    }

    @Override
    protected BstNode<K, V> add(BstNode<K, V> parent, BstNode<K, V> child) {
        if(parent == null){
            return child;
        }

        int flag = child.compareTo(parent);
        if(flag > 0){
            parent.right = add(parent.right, child);
        }else if(flag < 0){
            parent.left = add(parent.left, child);
        }else{
            // KEY相同，更新Value的值
            parent.value = child.value;
        }

        // 维护节点平衡
        parent = balanceNode(parent);

        return parent;
    }

    @Override
    protected BstNode<K, V> remove(BstNode<K, V> node, K key) {
        if(node == null){
            return null;
        }

        BstNode retNode = null;
        int flag = key.compareTo(node.key);
        if(flag < 0){
            node.left = remove(node.left, key);
            retNode = node;
        }else if(flag > 0){
            node.right = remove(node.right, key);
            retNode = node;
        }else{
            if(node.left == null){
                // 该节点的左子树为空
                BstNode right = node.right;
                node.right = null;
                retNode = right;
            }else if(node.right == null){
                // 该节点的右子树为空
                BstNode left = node.left;
                node.left = null;
                retNode = left;
            }else{
                // 左右子树均不为空
                // 找到比待删除的节点大的最小节点，即右子树中的最小节点
                BstNode<K, V> target = minimum(node.right);
                target.right = remove(node.right, target.key);
                target.left = node.left;
                node.left = node.right = null;
                retNode = target; // 将目标节点返回 给上级
            }
        }

        // 维护节点平衡
        retNode = balanceNode(retNode);

        return retNode;
    }
}

